#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <io.h>
#include <tchar.h>
#define PATH_SEP '\\'
#define HELPERS_DIR "C:\\arch-test\\"  // Windows默认辅助程序路径（可修改）
#define MAX_PATH_LEN MAX_PATH
#else
#include <fcntl.h>
#include <utime.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <dirent.h>
#include <signal.h>
#include <sys/wait.h>
#include <sys/resource.h>
#include <limits.h>
#define PATH_SEP '/'
#define HELPERS_DIR "/usr/lib/arch-test/"  // 类Unix默认辅助程序路径
#define MAX_PATH_LEN PATH_MAX
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

// 全局变量：保存chroot环境下的临时辅助程序路径（用于信号清理）
char *temp_helper_path = NULL;

/**
 * 信号/退出处理函数：清理临时文件
 */
#ifdef _WIN32
BOOL WINAPI signal_handler(DWORD dwCtrlType) {
#else
void signal_handler(int sig) {
#endif
    if (temp_helper_path != NULL && temp_helper_path[0] != '\0') {
#ifdef _WIN32
        DeleteFileA(temp_helper_path);
#else
        unlink(temp_helper_path);
#endif
        free(temp_helper_path);
        temp_helper_path = NULL;
    }

#ifdef _WIN32
    return TRUE;
#else
    // 重新注册信号（部分信号触发后会恢复默认处理）
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
#ifdef SIGEXIT
    signal(SIGEXIT, signal_handler);
#endif
#endif
}

/**
 * 设置禁止核心转储
 */
int set_no_core_dump(void) {
#ifdef _WIN32
    // Windows：禁用崩溃弹窗并设置错误模式
    SetErrorMode(SEM_NOGPFAULTERRORBOX | SEM_FAILCRITICALERRORS);
    return 0;
#else
    // 类Unix：通过rlimit限制核心转储大小为0
    struct rlimit rlim;
    rlim.rlim_cur = 0;
    rlim.rlim_max = 0;
    if (setrlimit(RLIMIT_CORE, &rlim) == -1) {
        perror("设置禁止核心转储失败");
        return -1;
    }
    return 0;
#endif
}

/**
 * 复制文件（保留权限和时间属性）
 * @param src 源文件路径
 * @param dest 目标文件路径
 * @return 0成功，-1失败
 */
int copy_file(const char *src, const char *dest) {
#ifdef _WIN32
    // Windows：使用系统API复制文件（保留属性）
    if (!CopyFileA(src, dest, FALSE)) {
        DWORD err = GetLastError();
        char err_msg[256];
        FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                      NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      err_msg, sizeof(err_msg), NULL);
        fprintf(stderr, "复制文件失败：%s（源：%s，目标：%s）\n", err_msg, src, dest);
        return -1;
    }
    return 0;
#else
    // 类Unix：手动实现文件复制（内容+权限+时间）
    int src_fd, dest_fd;
    struct stat src_stat;
    char buf[4096];
    ssize_t n;

    // 打开源文件
    if ((src_fd = open(src, O_RDONLY)) == -1) {
        perror("打开源文件失败");
        return -1;
    }

    // 获取源文件属性
    if (fstat(src_fd, &src_stat) == -1) {
        perror("获取源文件属性失败");
        close(src_fd);
        return -1;
    }

    // 创建目标文件（继承源文件权限）
    if ((dest_fd = open(dest, O_WRONLY | O_CREAT | O_TRUNC, src_stat.st_mode & 0777)) == -1) {
        perror("创建目标文件失败");
        close(src_fd);
        return -1;
    }

    // 复制文件内容
    while ((n = read(src_fd, buf, sizeof(buf))) > 0) {
        if (write(dest_fd, buf, n) != n) {
            perror("写入目标文件失败");
            close(src_fd);
            close(dest_fd);
            unlink(dest);  // 删除不完整文件
            return -1;
        }
    }

    if (n == -1) {
        perror("读取源文件失败");
        close(src_fd);
        close(dest_fd);
        unlink(dest);
        return -1;
    }

    // 关闭文件描述符
    close(src_fd);
    close(dest_fd);

    // 复制访问/修改时间
    struct utimbuf times = {src_stat.st_atime, src_stat.st_mtime};
    if (utime(dest, &times) == -1) {
        perror("设置目标文件时间失败（非致命）");
    }

    return 0;
#endif
}

/**
 * 执行辅助程序并获取输出
 * @param cmd 程序路径
 * @param output_buf 输出缓冲区
 * @param buf_size 缓冲区大小
 * @return 程序退出码（-1表示执行失败）
 */
int execute_helper(const char *cmd, char *output_buf, size_t buf_size) {
    if (output_buf == NULL || buf_size == 0) {
        fprintf(stderr, "输出缓冲区无效\n");
        return -1;
    }
    memset(output_buf, 0, buf_size);

#ifdef _WIN32
    // Windows：使用CreateProcess创建进程并通过管道获取输出
    HANDLE hReadPipe, hWritePipe;
    SECURITY_ATTRIBUTES sa = {sizeof(SECURITY_ATTRIBUTES), NULL, TRUE};
    PROCESS_INFORMATION pi;
    STARTUPINFOA si = {0};
    DWORD bytes_read;

    // 创建管道
    if (!CreatePipe(&hReadPipe, &hWritePipe, &sa, 0)) {
        DWORD err = GetLastError();
        char err_msg[256];
        FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                      NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      err_msg, sizeof(err_msg), NULL);
        fprintf(stderr, "创建管道失败：%s\n", err_msg);
        return -1;
    }

    // 配置子进程启动信息（重定向stdout/stderr）
    si.cb = sizeof(STARTUPINFOA);
    si.hStdOutput = hWritePipe;
    si.hStdError = hWritePipe;
    si.dwFlags |= STARTF_USESTDHANDLES;

    // 创建子进程
    if (!CreateProcessA(NULL, (LPSTR)cmd, NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) {
        DWORD err = GetLastError();
        char err_msg[256];
        FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                      NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      err_msg, sizeof(err_msg), NULL);
        fprintf(stderr, "启动进程失败（命令：%s）：%s\n", cmd, err_msg);
        CloseHandle(hReadPipe);
        CloseHandle(hWritePipe);
        return -1;
    }

    // 关闭写端管道（子进程已继承）
    CloseHandle(hWritePipe);

    // 读取子进程输出
    if (!ReadFile(hReadPipe, output_buf, buf_size - 1, &bytes_read, NULL)) {
        DWORD err = GetLastError();
        char err_msg[256];
        FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                      NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      err_msg, sizeof(err_msg), NULL);
        fprintf(stderr, "读取进程输出失败：%s\n", err_msg);
    } else {
        // 去除换行符和回车符（与原脚本行为一致）
        char *nl = strchr(output_buf, '\n');
        if (nl) *nl = '\0';
        char *cr = strchr(output_buf, '\r');
        if (cr) *cr = '\0';
    }

    // 清理资源
    CloseHandle(hReadPipe);
    WaitForSingleObject(pi.hProcess, INFINITE);
    
    DWORD exit_code = -1;
    GetExitCodeProcess(pi.hProcess, &exit_code);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);

    return (int)exit_code;
#else
    // 类Unix：使用fork+exec+pipe实现进程创建和输出捕获
    int pipefd[2];
    pid_t pid;
    ssize_t n;
    int status;

    // 创建管道
    if (pipe(pipefd) == -1) {
        perror("创建管道失败");
        return -1;
    }

    // 创建子进程
    if ((pid = fork()) == -1) {
        perror("创建子进程失败");
        close(pipefd[0]);
        close(pipefd[1]);
        return -1;
    }

    if (pid == 0) {  // 子进程：执行目标程序
        close(pipefd[0]);  // 关闭读端
        // 重定向stdout/stderr到管道写端
        if (dup2(pipefd[1], STDOUT_FILENO) == -1 || dup2(pipefd[1], STDERR_FILENO) == -1) {
            perror("重定向标准输出失败");
            exit(1);
        }
        close(pipefd[1]);  // 关闭原写端

        // 执行程序（参数数组形式）
        char *argv[] = {(char *)cmd, NULL};
        execvp(cmd, argv);
        perror("执行程序失败");  // exec失败才会执行到这里
        exit(1);
    }

    // 父进程：读取输出并等待子进程结束
    close(pipefd[1]);  // 关闭写端
    n = read(pipefd[0], output_buf, buf_size - 1);
    if (n > 0) {
        output_buf[n] = '\0';
        // 去除换行符和回车符
        char *nl = strchr(output_buf, '\n');
        if (nl) *nl = '\0';
        char *cr = strchr(output_buf, '\r');
        if (cr) *cr = '\0';
    } else if (n == -1) {
        perror("读取管道输出失败");
    }
    close(pipefd[0]);

    // 等待子进程结束并获取退出码
    if (waitpid(pid, &status, 0) == -1) {
        perror("等待子进程失败");
        return -1;
    }
    return WIFEXITED(status) ? WEXITSTATUS(status) : -1;
#endif
}

/**
 * 解析命令行选项
 * @param argc 参数个数
 * @param argv 参数数组
 * @param native_flag 输出：是否启用--native模式
 * @param chroot_dir 输出：chroot路径（需手动free）
 * @return 架构参数（需手动free），NULL表示无架构参数
 */
char *parse_options(int argc, char *argv[], int *native_flag, char **chroot_dir) {
    *native_flag = 0;
    *chroot_dir = NULL;
    char *arch = NULL;
    int i = 1;

    while (i < argc) {
        if (strcmp(argv[i], "--version") == 0) {
            printf("arch-test v0.20\n");
            exit(0);
        } else if (strcmp(argv[i], "-n") == 0 || strcmp(argv[i], "--native") == 0) {
            *native_flag = 1;
            // 设置环境变量（与原脚本一致）
#ifdef _WIN32
            SetEnvironmentVariableA("QEMU_VERSION", "meow");
            SetEnvironmentVariableA("WINEPREFIX", "\\dev\\null");
#else
            setenv("QEMU_VERSION", "meow", 1);
            setenv("WINEPREFIX", "/dev/null", 1);
#endif
            i++;
        } else if (strcmp(argv[i], "-c") == 0 || strcmp(argv[i], "--chroot") == 0) {
            i++;
            if (i >= argc) {
                fprintf(stderr, "%s: -c 需要指定chroot路径（可使用'.'表示当前目录）\n", argv[0]);
                exit(2);
            }

            // 分配并保存chroot路径
            *chroot_dir = strdup(argv[i]);
            if (*chroot_dir == NULL) {
                perror("内存分配失败");
                exit(1);
            }

#ifdef _WIN32
            // Windows不支持chroot功能
            fprintf(stderr, "%s: Windows系统不支持chroot功能\n", argv[0]);
            free(*chroot_dir);
            exit(2);
#else
            // 检查chroot路径有效性
            struct stat st;
            if (stat(*chroot_dir, &st) == -1) {
                perror("检查chroot路径失败");
                free(*chroot_dir);
                exit(2);
            }
            if (!S_ISDIR(st.st_mode)) {
                fprintf(stderr, "%s: chroot路径'%s'不是有效目录\n", argv[0], *chroot_dir);
                free(*chroot_dir);
                exit(2);
            }

            // 检查root权限
            if (getuid() != 0) {
                fprintf(stderr, "%s: 执行chroot需要root权限\n", argv[0]);
                free(*chroot_dir);
                exit(2);
            }
#endif

            i++;
            // 检查chroot模式下是否有且仅有一个架构参数
            if (i >= argc) {
                fprintf(stderr, "%s: -c 必须配合单个架构参数使用\n", argv[0]);
                free(*chroot_dir);
                exit(2);
            }
            if (arch != NULL) {
                fprintf(stderr, "%s: -c 模式下仅支持单个架构查询\n", argv[0]);
                free(*chroot_dir);
                exit(2);
            }

            // 保存架构参数
            arch = strdup(argv[i]);
            if (arch == NULL) {
                perror("内存分配失败");
                free(*chroot_dir);
                exit(1);
            }
            i++;
        } else if (argv[i][0] == '-') {
            fprintf(stderr, "%s: 未知选项'%s'\n", argv[0], argv[i]);
            exit(2);
        } else {
            // 处理架构参数（非选项参数）
            if (arch != NULL) {
                fprintf(stderr, "用法：%s [-n] 或 %s [-n|-c <chroot>] <架构>\n", argv[0], argv[0]);
                exit(1);
            }
            arch = strdup(argv[i]);
            if (arch == NULL) {
                perror("内存分配失败");
                exit(1);
            }
            i++;
        }
    }

    return arch;
}

/**
 * 列出所有支持的架构（遍历辅助程序目录）
 * @param helpers_dir 辅助程序目录路径
 * @return 0成功，-1失败
 */
int list_supported_archs(const char *helpers_dir) {
#ifdef _WIN32
    // Windows：使用FindFirstFile/FindNextFile遍历目录
    WIN32_FIND_DATAA find_data;
    HANDLE h_find;
    char search_path[MAX_PATH_LEN];
    char helper_path[MAX_PATH_LEN];
    char output[256];

    // 构建搜索路径（目录下所有文件）
    snprintf(search_path, sizeof(search_path), "%s*", helpers_dir);
    h_find = FindFirstFileA(search_path, &find_data);

    if (h_find == INVALID_HANDLE_VALUE) {
        DWORD err = GetLastError();
        if (err != ERROR_FILE_NOT_FOUND) {
            char err_msg[256];
            FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                          NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                          err_msg, sizeof(err_msg), NULL);
            fprintf(stderr, "遍历目录失败：%s\n", err_msg);
            return -1;
        }
        return 0;  // 目录为空，无输出
    }

    // 遍历目录项
    do {
        // 跳过.和..目录
        if (strcmp(find_data.cFileName, ".") == 0 || strcmp(find_data.cFileName, "..") == 0) {
            continue;
        }

        // 仅处理文件（跳过子目录）
        if (find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
            continue;
        }

        // 构建辅助程序完整路径
        snprintf(helper_path, sizeof(helper_path), "%s%s", helpers_dir, find_data.cFileName);

        // 提取架构名（去除"arch-test-"前缀和.exe后缀）
        char arch[256];
        const char *prefix = "arch-test-";
        size_t prefix_len = strlen(prefix);
        const char *name = find_data.cFileName;

        // 处理前缀
        if (strncmp(name, prefix, prefix_len) == 0) {
            name += prefix_len;
        }

        // 处理.exe后缀
        const char *ext = strrchr(name, '.');
        if (ext != NULL && _stricmp(ext, ".exe") == 0) {
            snprintf(arch, sizeof(arch), "%.*s", (int)(ext - name), name);
        } else {
            strncpy(arch, name, sizeof(arch) - 1);
            arch[sizeof(arch) - 1] = '\0';
        }

        // 执行辅助程序并判断是否支持
        int exit_code = execute_helper(helper_path, output, sizeof(output));
        if (exit_code == 0 && strcmp(output, "ok") == 0) {
            printf("%s\n", arch);
        }
    } while (FindNextFileA(h_find, &find_data));

    // 检查遍历是否正常结束
    DWORD err = GetLastError();
    if (err != ERROR_NO_MORE_FILES) {
        char err_msg[256];
        FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                      NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      err_msg, sizeof(err_msg), NULL);
        fprintf(stderr, "遍历目录异常：%s\n", err_msg);
        FindClose(h_find);
        return -1;
    }

    FindClose(h_find);
#else
    // 类Unix：使用opendir/readdir/closedir遍历目录
    DIR *dir = opendir(helpers_dir);
    if (dir == NULL) {
        perror("打开辅助程序目录失败");
        return -1;
    }

    struct dirent *entry;
    char helper_path[MAX_PATH_LEN];
    char output[256];

    // 遍历目录项
    while ((entry = readdir(dir)) != NULL) {
        // 跳过.和..目录
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) {
            continue;
        }

        // 构建辅助程序完整路径
        snprintf(helper_path, sizeof(helper_path), "%s%c%s", 
                 helpers_dir, PATH_SEP, entry->d_name);

        // 检查是否为可执行文件
        struct stat st;
        if (stat(helper_path, &st) == -1 || !S_ISREG(st.st_mode) || !(st.st_mode & S_IXUSR)) {
            continue;
        }

        // 提取架构名（去除"arch-test-"前缀）
        char arch[256];
        const char *prefix = "arch-test-";
        size_t prefix_len = strlen(prefix);
        if (strncmp(entry->d_name, prefix, prefix_len) == 0) {
            strncpy(arch, entry->d_name + prefix_len, sizeof(arch) - 1);
        } else {
            strncpy(arch, entry->d_name, sizeof(arch) - 1);
        }
        arch[sizeof(arch) - 1] = '\0';

        // 执行辅助程序并判断是否支持
        int exit_code = execute_helper(helper_path, output, sizeof(output));
        if (exit_code == 0 && strcmp(output, "ok") == 0) {
            printf("%s\n", arch);
        }
    }

    closedir(dir);
#endif

    return 0;
}

int main(int argc, char *argv[]) {
    int native_flag;
    char *chroot_dir = NULL;
    char *arch = NULL;
    char helper_path[MAX_PATH_LEN];
    char output[256];

    // 初始化：禁止核心转储
    if (set_no_core_dump() != 0) {
        fprintf(stderr, "警告：无法禁用核心转储\n");
    }

    // 注册信号处理（清理临时文件）
#ifdef _WIN32
    SetConsoleCtrlHandler(signal_handler, TRUE);
#else
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
#ifdef SIGEXIT
    signal(SIGEXIT, signal_handler);
#endif
#endif

    // 解析命令行选项
    arch = parse_options(argc, argv, &native_flag, &chroot_dir);

    // 处理架构查询（指定了架构参数）
    if (arch != NULL) {
        // 构建辅助程序路径
#ifdef _WIN32
        snprintf(helper_path, sizeof(helper_path), "%s%s.exe", HELPERS_DIR, arch);
#else
        snprintf(helper_path, sizeof(helper_path), "%s%c%s", 
                 HELPERS_DIR, PATH_SEP, arch);
#endif

        // 检查辅助程序是否存在且可执行
#ifdef _WIN32
        DWORD attr = GetFileAttributesA(helper_path);
        if (attr == INVALID_FILE_ATTRIBUTES || (attr & FILE_ATTRIBUTE_DIRECTORY)) {
            fprintf(stderr, "我不知道如何检测架构'%s'，抱歉\n", arch);
            goto cleanup;
        }
        // Windows下检查.exe后缀（简化可执行性判断）
        const char *ext = strrchr(helper_path, '.');
        if (ext == NULL || _stricmp(ext, ".exe") != 0) {
            fprintf(stderr, "我不知道如何检测架构'%s'，抱歉\n", arch);
            goto cleanup;
        }
#else
        struct stat st;
        if (stat(helper_path, &st) == -1 || !S_ISREG(st.st_mode) || !(st.st_mode & S_IXUSR)) {
            fprintf(stderr, "我不知道如何检测架构'%s'，抱歉\n", arch);
            goto cleanup;
        }
#endif

        int exit_code;

#ifndef _WIN32
        // 处理chroot模式（仅类Unix支持）
        if (chroot_dir != NULL) {
            // 提取辅助程序文件名
            const char *helper_name = strrchr(helper_path, PATH_SEP);
            helper_name = helper_name ? helper_name + 1 : helper_path;

            // 构建chroot下的临时路径
            temp_helper_path = malloc(MAX_PATH_LEN);
            if (temp_helper_path == NULL) {
                perror("内存分配失败");
                goto cleanup;
            }
            snprintf(temp_helper_path, MAX_PATH_LEN, "%s%c%s", 
                     chroot_dir, PATH_SEP, helper_name);

            // 复制辅助程序到chroot目录
            if (copy_file(helper_path, temp_helper_path) != 0) {
                fprintf(stderr, "复制辅助程序到chroot目录失败\n");
                goto cleanup;
            }

            // 切换到chroot环境
            if (chroot(chroot_dir) == -1) {
                perror("执行chroot失败");
                goto cleanup;
            }
            if (chdir("/") == -1) {
                perror("切换到chroot根目录失败");
                goto cleanup;
            }

            // 执行chroot内的辅助程序
            char chroot_helper[256];
            snprintf(chroot_helper, sizeof(chroot_helper), "/%s", helper_name);
            exit_code = execute_helper(chroot_helper, output, sizeof(output));
        } else {
#endif
            // 正常模式：直接执行辅助程序
            exit_code = execute_helper(helper_path, output, sizeof(output));
#ifndef _WIN32
        }
#endif

        // 输出检测结果
        if (exit_code == 0 && strcmp(output, "ok") == 0) {
            printf("%s: 支持\n", arch);
            goto cleanup_success;
        } else {
            printf("%s: 此机器/内核不支持\n", arch);
            goto cleanup;
        }
    }

    // 无架构参数：列出所有支持的架构
    list_supported_archs(HELPERS_DIR);

cleanup_success:
    signal_handler(0);  // 清理临时文件
    free(arch);
    free(chroot_dir);
    return 0;

cleanup:
    signal_handler(0);  // 清理临时文件
    free(arch);
    free(chroot_dir);
    return 2;
}

